import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from net import AlexNet


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model
model = AlexNet(num_classes =10, in_channels =3)
model.load_state_dict(torch.load('best_alexnet_model.pth'))
model.to(device)
model.eval()

# 测试时增强（Test-Time Augmentation, TTA）
# 在测试时，网络通过提取5个224x224块（四个边角块和一个中心块）
# 以及它们的水平翻转（因此共十个块）做预测，
# 然后网络的softmax层对这十个块做出的预测取均值。

# 数据预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.RandomCrop(227), #TTA 时内部做了，注释掉
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.4595, 0.4402, 0.3700], std=[0.2005, 0.1984, 0.1907])
])
# 测试数据
test_dataset = datasets.ImageFolder(root='F:/深度学习数据集/imagenet-10/test/', transform=transform)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4)

def tta_predict(model, image):
    """
    对单张图像进行TTA预测
    """
    # 裁剪大小
    crop_size = 227
    # 初始化预测结果
    predictions = []

    # 四个角和中心的裁剪
    crops = [
        image[:, :crop_size, :crop_size],  # 左上角
        image[:, :crop_size, -crop_size:],  # 右上角
        image[:, -crop_size:, :crop_size],  # 左下角
        image[:, -crop_size:, -crop_size:],  # 右下角
        transforms.CenterCrop(crop_size)(image)  # 中心裁剪
    ]

    # 对每个裁剪块进行预测
    for crop in crops:
        # 原始块
        predictions.append(model(crop.unsqueeze(0)).softmax(dim=1))
        # 水平翻转块
        flipped_crop = crop.flip(dims=[2])  # 水平翻转
        predictions.append(model(flipped_crop.unsqueeze(0)).softmax(dim=1))

    # 取均值
    avg_prediction = torch.stack(predictions).mean(dim=0)
    return avg_prediction

if __name__ == "__main__":

    correct = 0
    total = 0
    with torch.no_grad():
        for data, label in test_loader:
            data, label = data.to(device), label.to(device)

            # Accuracy 78.08333333333333%
            # output = tta_predict(model, data.squeeze(0))  # data.squeeze(0) 去掉batch维度,TTA操作针对单张，不是batch

            # Accuracy 73.66666666666667%
            output = model(data)
            _, predicted = torch.max(output.data, 1)# 预测的类别
            total += label.size(0)
            correct += (predicted == label).sum().item()

    print(f'Accuracy of the network on the test images: {100 * correct / total}%')
